package com.kjhxtc.restfulrunner;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.StrUtil;
import lombok.SneakyThrows;
import org.apache.commons.cli.*;

import java.io.ByteArrayOutputStream;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CountDownLatch;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @author 杜神
 * @version 1.0
 * @since 2022/3/16 9:30
 */
public class Runner extends Thread {
    public static String HEADER = "Restful API压测工具V1.0.0";
    public static String FOOTER = "Create by 杜神 at 2022-03-16";
    private static final Options OPTIONS = new Options();
    private static String HELP_STRING = null;
    public static Integer THREAD_NUM = 8;
    public static Integer TIMES = 1000;
    public static long DURATION = 0;
    public static int MAX_TOTAL = 20;
    public static int MAX_PER_ROUTE = 50;
    public static boolean KEEP_ALIVE;
    private final CountDownLatch countDown;
    private final TaskManager taskManager;
    public static String HOST;
    public static Integer PORT;

    public Runner(CountDownLatch downLatch, String name, TaskManager taskManager) {
        this.setName(name);
        this.countDown = downLatch;
        this.taskManager = taskManager;
    }

    @Override
    public void run() {
        while (!taskManager.isEmpty()) {
            Task task = taskManager.pop();
            Log.trace("{} 线程号:{},task:{}", DateUtil.now(), getName(), task.hashCode());
            task.execute();
        }
        countDown.countDown();
    }

    @SneakyThrows
    public static void main(String[] args) {
        CommandLineParser commandLineParser = new DefaultParser();
        OPTIONS.addOption(Option.builder("f").argName("保存http请求的文件路径").desc("必填项").hasArg(true).longOpt("file").required().type(String.class).build());
        OPTIONS.addOption(Option.builder("t").argName("线程数").desc("默认:" + THREAD_NUM).hasArg(true).longOpt("threads").type(Integer.class).build());
        OPTIONS.addOption(Option.builder("m").argName("请求次数或时间,").hasArg(true).longOpt("mode").type(String.class).desc("e.g. 5000(5000次),10s(10秒),1m(1分钟),1h(1小时),1d(1天)\n默认:" + TIMES).build());
        OPTIONS.addOption(Option.builder("mt").argName("客户端总并行链接最大数").desc("默认:" + MAX_TOTAL).hasArg(true).longOpt("max_total").type(Integer.class).build());
        OPTIONS.addOption(Option.builder("mp").argName("每个主机的最大并行链接数").desc("默认:" + MAX_PER_ROUTE).hasArg(true).longOpt("max_per_route").type(Integer.class).build());
        OPTIONS.addOption(Option.builder("a").desc("连接是否复用，默认:" + KEEP_ALIVE).longOpt("alive").build());
        OPTIONS.addOption(Option.builder("v").desc("info输出级别").longOpt("info").build());
        OPTIONS.addOption(Option.builder("vv").desc("debug输出级别").longOpt("debug").build());
        OPTIONS.addOption(Option.builder("vvv").desc("trace输出级别").longOpt("trace").build());
        OPTIONS.addOption(Option.builder("help").desc("帮助").longOpt("help").build());
        String command = null;
        try {
            CommandLine commandLine = commandLineParser.parse(OPTIONS, args);
            Log.IS_INFO = commandLine.hasOption("v") || commandLine.hasOption("vv") || commandLine.hasOption("vvv");
            Log.IS_DEBUG = commandLine.hasOption("vv") || commandLine.hasOption("vvv");
            Log.IS_TRACE = commandLine.hasOption("vvv");
            if (commandLine.hasOption("help")) {
                Log.log(getHelpString());
            }
            if (commandLine.hasOption("f")) {
                String filePath = commandLine.getOptionValue("f");
                command = StrUtil.trim(FileUtil.readString(filePath, StandardCharsets.UTF_8));
                String[] groups = command.split("\n");
                Assert.notEmpty(groups, "请求文件解析错误");
                Pattern pattern = Pattern.compile("^\\s*[a-zA-Z]+\\s+(https?)://([^/]+).*\\s*$");
                Matcher matcher = pattern.matcher(groups[0]);
                Assert.isTrue(matcher.matches(), "HOST解析失败:{}",groups[0]);
                String host = matcher.group(2);
                Assert.notEmpty(host, "解析host错误");
                if (!host.contains(":")) {
                    PORT = StrUtil.equalsAnyIgnoreCase("https", matcher.group(1)) ? 443 : 80;
                    HOST = host;
                } else {
                    String[] split = host.split(":");
                    Assert.isTrue(split.length == 2, "解析错误:{}", host);
                    HOST = split[0];
                    PORT = Integer.parseInt(split[1]);
                }
            }
            if (commandLine.hasOption("t")) {
                THREAD_NUM = Integer.parseInt(commandLine.getOptionValue("t"));
            }
            if (commandLine.hasOption("H")) {
                HOST = commandLine.getOptionValue("H");
            }
            if (commandLine.hasOption("P")) {
                PORT = Integer.parseInt(commandLine.getOptionValue("P"));
            }
            if (commandLine.hasOption("mt")) {
                MAX_TOTAL = Integer.parseInt(commandLine.getOptionValue("mt"));
            }
            if (commandLine.hasOption("mp")) {
                MAX_PER_ROUTE = Integer.parseInt(commandLine.getOptionValue("mp"));
            }
            if (commandLine.hasOption("a")) {
                KEEP_ALIVE = true;
            }
            if (commandLine.hasOption("m")) {
                String m = commandLine.getOptionValue("m");
                if (Pattern.matches("^\\d+[smhSd]$", m)) {
                    long a = Long.parseLong(m.substring(0, m.length() - 1));
                    String b = m.substring(m.length() - 1);
                    if ("s".equals(b)) {
                        a *= 1000;
                    }
                    if ("m".equals(b)) {
                        a *= 1000 * 60;
                    }
                    if ("h".equals(b)) {
                        a *= 1000 * 60 * 60;
                    }
                    if ("d".equals(b)) {
                        a *= 1000 * 60 * 60 * 24;
                    }
                    DURATION = a;
                } else if (Pattern.matches("^\\d+$", m)) {
                    TIMES = Integer.parseInt(m);
                }
            }
        } catch (ParseException e) {
            Log.log(e.getMessage() + "\n" + getHelpString());
            System.exit(0);
        }

        CountDownLatch countDownLatch = new CountDownLatch(THREAD_NUM);
        RestfulClient client = new RestfulClient(HOST, PORT, MAX_TOTAL, MAX_PER_ROUTE);
        RestfulClientManager clientManager = new RestfulClientManager();
        clientManager.add(client);
        ResultCollector resultCollector = new ResultCollector();
        TaskManager tasks;
        Task task = new Task(command, clientManager, resultCollector);
        if (DURATION == 0) {
            tasks = new TaskManager();
            for (int i = 0; i < TIMES; i++) {
                tasks.push(task);
            }
        } else {
            tasks = new TaskManager(System.currentTimeMillis() + DURATION);
            tasks.push(task);
        }
        long beginTime = System.currentTimeMillis();
        for (int i = 0; i < THREAD_NUM; i++) {
            Runner runner = new Runner(countDownLatch, "t" + i, tasks);
            runner.start();
        }
        try {
            countDownLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        long endTime = System.currentTimeMillis();
        Log.log("\nfinish!");
        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        Log.log("-----------------{}--------------------", HEADER);
        Log.log(FOOTER);
        Log.log("now {}", DateUtil.now());
        Log.log("------------------------------------------------------------");
        // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        Log.log("\n>>>>>>>>压测参数");
        Log.log("UV:{}\nMAX_TOTAL:{}\nMAX_PER_ROUTE:{}\nDURATION:{}ms\nSAMPLES:{}\nalive:{}" +
                "\nHOST:{}" +
                "\nPORT:{}", THREAD_NUM, MAX_TOTAL, MAX_PER_ROUTE, DURATION, DURATION == 0 ? TIMES : "N/A", KEEP_ALIVE ? "长连接" : "短连接", HOST, PORT);
        Log.log("Request >>");
        Log.log(command);
        Log.log("<<");

        Log.log("\n>>>>>>>>结果统计");
        long costTime = endTime - beginTime;
        Log.log("总耗时:{}", DateUtil.formatBetween(costTime));
        Log.log("交易总数:{}", resultCollector.count());
//        Log.log("QPS:" + resultCollector.qps(costTime));
        Log.log("TPS:{}", resultCollector.qps(costTime));
        Log.log("平均响应时间:{}ms", resultCollector.avgCost());
        Log.log("响应时间中位数:{}ms", resultCollector.midCost());
        Log.log("成功率:{}%", resultCollector.successRate());
    }

    private static String getHelpString() {
        if (HELP_STRING == null) {
            HelpFormatter helpFormatter = new HelpFormatter();
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            PrintWriter printWriter = new PrintWriter(byteArrayOutputStream);
            helpFormatter.printHelp(printWriter, HelpFormatter.DEFAULT_WIDTH, " -help", HEADER + "\n" + FOOTER,
                    OPTIONS, HelpFormatter.DEFAULT_LEFT_PAD, HelpFormatter.DEFAULT_DESC_PAD, "\n请求文件格式：\nGET http://127.0.0.1:8080/test\n" +
                            "Content-Type:application/json\n" +
                            "\n" +
                            "{\n" +
                            "\"key\":1\n" +
                            "}");
            printWriter.flush();
            HELP_STRING = byteArrayOutputStream.toString();
            printWriter.close();
        }
        return HELP_STRING;
    }
}
